package com.bigdata.hive;

import cn.hutool.db.Db;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.security.UserGroupInformation;

import java.io.IOException;
import java.sql.*;
import java.util.*;

/**
 * @author xin.ding
 * @date 2020/10/16 13:51
 */
public class HiveUtils {
    private static final String DRIVER = "org.apache.hive.jdbc.HiveDriver";

    /**
     * krb5.conf路径
     */
    private String krb5Path;

    /**
     * keytab文件路径
     */
    private String keytabPath;

    /**
     * hive仓库的连接地址
     */
    private String url;

    /**
     * hive对应用户
     */
    private String user;

    /**
     * hive用户对应的密码
     */
    private String pwd;

    private boolean isKerberos;

    /**
     * Hive 连接
     */
    private static Connection conn = null;

    public HiveUtils(String krb5Path, String keytabPath, String url, String user, String pwd) {
        this.krb5Path = krb5Path;
        this.keytabPath = keytabPath;
        this.url = url;
        this.user = user;
        this.pwd = pwd;
        Db.use();
    }

    public static void main(String[] args) {

        String url = "jdbc:hive2://tdh01:10000/;principal=hive/tdh01@HADOOP.COM";
        String krb5Path = "D:\\keytab\\krb5.conf";
        String user = "hive/tdh01@HADOOP.COM";
        String keytabPath = "D:\\keytab\\hive.keytab";

        HiveUtils oper = new HiveUtils(krb5Path, keytabPath, url, user, "");
        List<String> dbs = oper.showDatabases();
        for (String db : dbs) {
            System.out.println(db);
        }

        List<String> tbls = oper.listTables("default");
        for (String tbl : tbls) {
            System.out.println(tbl);
        }

//        String url1 = "jdbc:hive2://tdh02:10000";
//        HiveUtils nokerberos = new HiveUtils(krb5Path, keytabPath, url1, user, "");
//
//        Connection connection = nokerberos.getConnection();
//        List<String> dbs1 = nokerberos.showDatabases();
//        for (String db : dbs1) {
//            System.out.println(db);
//        }

    }


    /**
     * 获取hive连接
     */
    private Connection getConnection() {
        try {
            if (conn == null || conn.isClosed()) {
                if (isKerberos) {
                    initKerberosEnv(krb5Path, user, keytabPath);
                }
                Class.forName(DRIVER);
                conn = DriverManager.getConnection(url);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        return conn;
    }


    /**
     * 构建hive Kerberos 配置信息
     *
     * @param krb5Path   krb5.conf路径
     * @param user       keytab principal 名称
     * @param keytabPath keytab文件路径
     */
    public static void initKerberosEnv(String krb5Path, String user, String keytabPath) {
        try {
            System.setProperty("java.security.krb5.conf", krb5Path);
            Configuration configuration = new Configuration();
            configuration.set("hadoop.security.authentication", "Kerberos");
            UserGroupInformation.setConfiguration(configuration);
            UserGroupInformation.loginUserFromKeytab(user, keytabPath);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 关闭连接
     */
    public void close() {
        try {
            if (conn != null && !conn.isClosed()) {
                conn.close();
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    /**
     * 关闭Statement
     *
     * @param stmt
     */
    public void close(Statement stmt) {
        try {
            if (stmt != null) {
                stmt.close();
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    /**
     * 关闭PreparedStatement
     *
     * @param pst
     */
    public void close(PreparedStatement pst) {
        try {
            if (pst != null) {
                pst.close();
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    /**
     * 关闭ResultSet
     *
     * @param rs
     */
    public void close(ResultSet rs) {
        try {
            if (rs != null) {
                rs.close();
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    /**
     * 列出指定数据库下的所有表
     *
     * @param dbName
     * @return
     */
    public List<String> listTables(String dbName) {
        Statement stmt = null;
        ResultSet res = null;
        List<String> tables = new LinkedList<String>();
        try {
            stmt = getConnection().createStatement();
            if (dbName != null && dbName.trim().length() > 0) {
                stmt.execute("USE " + dbName);
            }
            res = stmt.executeQuery("SHOW TABLES");
            while (res.next()) {
                tables.add(res.getString(1));
            }
        } catch (SQLException e) {
            throw new HiveDBException(e);
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return tables;
    }

    /*
     * 获取数据库
     */
    public List<String> showDatabases() {
        Statement stmt = null;
        ResultSet res = null;
        List<String> tables = new LinkedList<String>();
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery("SHOW DATABASES");
            while (res.next()) {
                tables.add(res.getString(1));
            }
        } catch (SQLException e) {
            throw new HiveDBException(e);
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return tables;
    }

    /**
     * 执行非查询的sql语句，比如创建表，加载数据等等
     *
     * @param sql
     * @return
     */
    public boolean executeNonQuery(String sql) {
        Statement stmt = null;
        boolean result = true;
        try {
            stmt = getConnection().createStatement();
            stmt.execute(sql);
        } catch (SQLException e) {
            result = false;
            throw new HiveDBException(e);
        } finally {
            close(stmt);
            close();
        }
        return result;
    }

    /**
     * 使用Statement查询数据，返回ResultSet
     *
     * @param sql
     * @return
     */
    public ResultSet queryForResultSet(String sql) {
        Statement stmt = null;
        ResultSet res = null;
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery(sql);
        } catch (SQLException e) {
            throw new HiveDBException(e);
        } finally {
            close(stmt);
        }
        return res;
    }

    /**
     * 使用Statement查询数据，返回List集合，数据量比较小的时候用
     *
     * @param sql
     * @return
     */
    public List<Map<String, Object>> queryForList(String sql) {
        Statement stmt = null;
        ResultSet res = null;
        List<Map<String, Object>> list = null;
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery(sql);
            Map<String, Object> map = null;
            ResultSetMetaData rsmd = res.getMetaData();
            int rowCnt = rsmd.getColumnCount();
            list = new LinkedList<Map<String, Object>>();
            while (res.next()) {
                map = new LinkedHashMap<String, Object>(rowCnt);
                for (int i = 1; i <= rowCnt; i++) {
                    map.put(rsmd.getColumnName(i), res.getObject(i));
                }
                list.add(map);
            }
        } catch (SQLException e) {
            throw new HiveDBException(e);
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return list;
    }

    /**
     * 使用PreparedStatement查询数据，返回ResultSet
     *
     * @param sql
     * @param values
     * @return
     */
    public ResultSet queryForResultSet(String sql, String[] values) {
        PreparedStatement pst = null;
        ResultSet res = null;
        try {
            pst = getConnection().prepareStatement(sql);
            setValue(pst, values);
            res = pst.executeQuery();
        } catch (SQLException e) {
            throw new HiveDBException(e);
        } finally {
            close(pst);
        }
        return res;
    }

    /**
     * 使用PreparedStatement查询数据，返回List集合，数据量比较小的时候用
     *
     * @param sql
     * @param values
     * @return
     */
    public List<Map<String, Object>> queryForList(String sql, String[] values) {
        PreparedStatement pst = null;
        ResultSet res = null;
        List<Map<String, Object>> list = null;
        try {
            pst = getConnection().prepareStatement(sql);
            setValue(pst, values);
            res = pst.executeQuery();
            Map<String, Object> map = null;
            ResultSetMetaData rsmd = res.getMetaData();
            int rowCnt = rsmd.getColumnCount();
            list = new LinkedList<Map<String, Object>>();
            while (res.next()) {
                map = new LinkedHashMap<String, Object>(rowCnt);
                for (int i = 1; i <= rowCnt; i++) {
                    map.put(rsmd.getColumnName(i), res.getObject(i));
                }
                list.add(map);
            }
        } catch (SQLException e) {
            throw new HiveDBException(e);
        } finally {
            close(res);
            close(pst);
            close();
        }
        return list;
    }


    /**
     * 执行数据文件导入
     *
     * @param sql
     * @return
     */
    public boolean impBySql(String sql) {
        PreparedStatement pst = null;
        boolean flag = false;
        try {
            pst = getConnection().prepareStatement(sql);
            flag = pst.execute();

        } catch (SQLException e) {
            throw new HiveDBException(e);
        } finally {
            close(pst);
            close();
        }
        return flag;
    }

    private void setValue(PreparedStatement pst, String[] values) {
        try {
            if (values == null || values.length == 0) {
                return;
            }
            for (int i = 0; i < values.length; i++) {
                pst.setString(i + 1, values[i]);
            }
        } catch (SQLException e) {
            throw new HiveDBException(e);
        }
    }

    /**
     * 获取表所在的路径
     *
     * @param tblName
     * @return
     * @throws Exception
     */
    public String getHiveTblPath(String tblName) throws Exception {
        String result = "";
        Statement stmt = null;
        ResultSet res = null;
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery("desc extended " + tblName);
            while (res.next()) {
                if (res.getString(1).trim().equals("Detailed Table Information")) {
                    String content = res.getString(2).trim();
                    int start = content.indexOf("location:");
                    if (start == -1) {
                        continue;
                    }

                    String sub = content.substring(start);
                    int end = sub.indexOf(",");
                    if (end == -1) {
                        continue;
                    }

                    result = sub.substring("location:".length(), end);
                } else {
                    continue;
                }

                // String content = res.getString(1).trim();

            }
        } catch (SQLException e) {
            throw new Exception(e);
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return result;
    }

    /**
     * 获取表所在的路径
     *
     * @param tblName
     * @return
     * @throws Exception
     */
    public Map<String, String> getNewHiveTblPath(String tblName) throws Exception {
        String result = "";
        String field = "";
        Statement stmt = null;
        ResultSet res = null;
        Map<String, String> map = new HashMap<String, String>();
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery("desc extended " + tblName);
            while (res.next()) {
                if (res.getString(1).trim().equals("Detailed Table Information")) {
                    String content = res.getString(2).trim();
                    field = getField(content);
                    int start = content.indexOf("dragoncluster");
                    if (start == -1) {
                        continue;
                    }

                    String sub = content.substring(start);
                    int end = sub.indexOf(",");
                    if (end == -1) {
                        continue;
                    }

                    result = sub.substring("dragoncluster".length(), end);
                } else {
                    continue;
                }

                // String content = res.getString(1).trim();

            }
        } catch (SQLException e) {
            throw new Exception(e);
        } finally {
            close(res);
            close(stmt);
            close();
        }
        map.put("field", field);
        map.put("hdfsPath", result);
        return map;
    }

    public String getField(String content) throws Exception {
        int start = content.indexOf("field.delim=");
        if (start == -1) {
            return "\\\\001";
        } else {
            String sub = content.substring(start);
//          int end = sub.indexOf("}");
            return sub.substring("field.delim=".length(), "field.delim=".length() + 1);

        }
    }


    public String getTblComment(String tblName) {
        String result = null;
        Statement stmt = null;
        ResultSet res = null;
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery("desc extended " + tblName);
            while (res.next()) {
                if (res.getString(1).trim().equals("Detailed Table Information")) {
                    String content = res.getString(2).trim();
                    int start = content.lastIndexOf("comment");
                    if (start == -1) {
                        continue;
                    }

                    String sub = content.substring(start);

                    int endBracket = sub.indexOf("}");
                    int endDot = sub.indexOf(",");
                    int end = endBracket < endDot ? endBracket : endDot;
                    if (end == -1) {
                        continue;
                    }

                    result = sub.substring("comment=".length(), end);
                    if (result != null && result.startsWith("null")) {
                        result = null;
                    }
                } else {
                    continue;
                }
            }
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return result;
    }

    /**
     * desc:获取hive表文件的类型
     *
     * @param tblName
     */
    public String getTblFileType(String tblName) {
        String result = null;
        Statement stmt = null;
        ResultSet res = null;
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery("desc extended " + tblName);
            while (res.next()) {
                if (res.getString(1).trim().equals("Detailed Table Information")) {
                    String content = res.getString(2).trim();
                    if (content.toUpperCase().contains("TEXTINPUTFORMAT")) {
                        result = "TEXTFILE";
                    } else if (content.toUpperCase().contains("SEQUENCEFILEINPUTFORMAT")) {
                        result = "SEQUENCEFILE";
                    } else {
                        result = "SEQUENCEFILE";
                    }
                }
            }
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return result;
    }

    /**
     * desc:查看表是否为外部表
     *
     * @param tblName
     */
    public boolean isExternalTbl(String tblName) {
        boolean result = false;
        Statement stmt = null;
        ResultSet res = null;
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery("desc extended " + tblName);
            while (res.next()) {
                if (res.getString(1).trim().equals("Detailed Table Information")) {
                    String content = res.getString(2).trim();
                    if (content.toUpperCase().contains("EXTERNAL_TABLE")) {
                        result = true;
                    }
                }
            }
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return result;
    }

    public List<String[]> getColumAndType(String tblName) {
        Statement stmt = null;
        ResultSet res = null;
        List<String[]> list = null;
        String[] item = null;
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery("desc formatted " + tblName);
            list = new LinkedList<String[]>();

            while (res.next()) {
                if (res.getString(1).trim().equals("# col_name")) {
                    continue;
                }

                if (res.getString(1).equals("# Detailed Table Information") || res.getString(1).equals("# Partition Information")) {
                    break;
                }

                if (res.getString(1).trim().equals("")) {
                    continue;
                }
                String column = res.getString(1).trim().toUpperCase();
                String type = res.getString(2).trim().toUpperCase();
                String comment = "";
                if (res.getString(3) != null && res.getString(3).trim().length() > 0) {
                    comment = res.getString(3).trim().toUpperCase();
                    if ("NONE".equals(comment)) {
                        comment = "";
                    }
                }
                item = new String[]{column, type, comment};

                list.add(item);
            }
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return list;

    }

    public List<String[]> getHiveTblPartitions(String tblName) {
        Statement stmt = null;
        ResultSet res = null;
        List<String[]> list = null;
        String[] item = null;
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery("desc " + tblName);
            list = new LinkedList<String[]>();

            while (res.next()) {
                if (res.getString(1).equals("# Partition Information")) {
                    while (res.next()) {
                        if (res.getString(1).trim().equals("# col_name")) {
                            continue;
                        }
                        if (res.getString(1).trim().equals("")) {
                            continue;
                        }
                        String column = res.getString(1).trim().toUpperCase();
                        String type = res.getString(2).trim().toUpperCase();
                        String comment = "";
                        if (res.getString(3) != null && res.getString(3).trim().length() > 0) {
                            comment = res.getString(3).trim().toUpperCase();
                            if ("NONE".equals(comment)) {
                                comment = "";
                            }
                        }
                        item = new String[]{column, type, comment};
                        list.add(item);
                    }
                }
            }
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return list;

    }

    public List<String> getHiveTblColumns(String tblName) {
        Statement stmt = null;
        ResultSet res = null;
        List<String> list = null;
        try {
            stmt = getConnection().createStatement();
            res = stmt.executeQuery("desc formatted " + tblName);
            list = new LinkedList<String>();
            while (res.next()) {
                if (res.getString(1).trim().equals("# col_name")) {
                    continue;
                }

                if (res.getString(1).equals("# Detailed Table Information") || res.getString(1).equals("# Partition Information")) {
                    break;
                }

                if (res.getString(1).trim().equals("")) {
                    continue;
                }
                System.out.println(res.getString(1).trim());
                list.add(res.getString(1).trim().toUpperCase());
            }
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return list;
    }

    /**
     * desc:查看某张表是否存在
     *
     * @param database
     * @param tableName
     * @return date:2016年10月25日
     * author:Tonny Chien
     */
    public boolean existTable(String database, String tableName) {
        boolean result = false;
        Statement stmt = null;
        ResultSet res = null;
        try {
            String hql = "SHOW TABLES IN " + database;
            stmt = getConnection().createStatement();
            res = stmt.executeQuery(hql);
            while (res.next()) {
                if (res.getString(1).trim().toUpperCase().equals(tableName.toUpperCase())) {
                    result = true;
                    break;
                }
            }
        } catch (Exception e) {
            result = false;
        } finally {
            close(res);
            close(stmt);
            close();
        }
        return result;
    }

    /**
     * @param cmdType  grant/revoke
     * @param privs    权限
     * @param tgtType  database/table
     * @param tgt      数据库名/表名
     * @param recvType group/user
     * @param recv
     * @return
     * @author Tonny Chien
     * @date 207-5-21 20:13
     */
    public boolean auth(AUTH cmdType, String privs, AUTH tgtType, String tgt, AUTH recvType, String recv) {
        // 拼接命令
        StringBuilder sb = new StringBuilder();
        switch (cmdType) {
            case grant:
                sb.append("GRANT ");
                break;
            case revoke:
                sb.append("REVOKE ");
                break;
            default:
                break;
        }

        sb.append(privs);
        sb.append(" ON ");

        switch (tgtType) {
            case database:
                sb.append("DATABASE ");
                break;
            case table:
                sb.append("TABLE ");
                break;
            default:
                break;
        }

        sb.append(tgt);

        switch (cmdType) {
            case grant:
                sb.append(" TO ");
                break;
            case revoke:
                sb.append(" FROM ");
                break;
            default:
                break;
        }

        switch (recvType) {
            case user:
                sb.append(" USER ");
                break;
            case group:
                sb.append(" GROUP ");
                break;
            default:
                break;
        }
        sb.append(recv);

        String hql = sb.toString();

        boolean result = false;
        Statement stmt = null;
        try {
            stmt = getConnection().createStatement();
            stmt.execute("set role admin");
            stmt.execute(hql);
            result = true;
        } catch (Exception e) {
            result = false;
            e.printStackTrace();
        } finally {
            close(stmt);
            close();
        }
        return result;
    }


    enum AUTH {
        grant,
        revoke,
        database,
        table,
        user,
        group
    }
}


class HiveDBException extends RuntimeException {
    private static final long serialVersionUID = 2637639405785985892L;

    public HiveDBException(Exception e) {
        super(e.getMessage());
    }

    public HiveDBException(String message, Throwable cause) {
        super(message, cause);
    }

    public HiveDBException(String message) {
        super(message);
    }

    public HiveDBException(Throwable cause) {
        super(cause);
    }
}